from PyQt5.QtWidgets import QOpenGLWidget, QSizePolicy
from PyQt5.QtCore import Qt
from OpenGL.GL import *
from OpenGL.GLUT import *
from OpenGL.GLU import *
import numpy as np
import open3d as o3d

class PointCloudViewer(QOpenGLWidget):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.point_cloud = None  # (M,3) 点云坐标
        self.point_colors = None  # (M,3) 点云颜色
        self.mesh = None          # Open3D 网格对象
        self.point_size = 2.0     # 点大小（默认2）
        self.opacity = 1.0        # 透明度（默认1.0）
        self.bg_color = (0.0, 0.0, 0.0)  # 背景色（默认黑色）
        self.show_axis = False    # 是否显示坐标轴
        self.rotation_x = 0
        self.rotation_y = 0
        self.zoom = 5.0
        self.last_mouse_pos = None

        self.setMinimumSize(300, 300)
        self.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)

    def initializeGL(self):
        glClearColor(*self.bg_color, 1.0)
        glEnable(GL_DEPTH_TEST)
        glEnable(GL_PROGRAM_POINT_SIZE)
        glEnable(GL_POINT_SMOOTH)
        glEnable(GL_LIGHTING)
        glEnable(GL_LIGHT0)
        glEnable(GL_COLOR_MATERIAL)
        glEnable(GL_NORMALIZE)

        # 光照参数（适配黑色背景）
        light_pos = [3.0, 3.0, 3.0, 0.0]
        glLightfv(GL_LIGHT0, GL_POSITION, light_pos)
        glLightfv(GL_LIGHT0, GL_AMBIENT, [0.5, 0.5, 0.5, 1.0])
        glLightfv(GL_LIGHT0, GL_DIFFUSE, [1.0, 1.0, 1.0, 1.0])
        glLightfv(GL_LIGHT0, GL_SPECULAR, [0.8, 0.8, 0.8, 1.0])

    def resizeGL(self, width, height):
        glViewport(0, 0, width, height)
        glMatrixMode(GL_PROJECTION)
        glLoadIdentity()
        gluPerspective(45, (width / height), 0.1, 100.0)
        glMatrixMode(GL_MODELVIEW)
        glLoadIdentity()

    def paintGL(self):
            # 每次绘制都重新设置背景色（关键修复）
        glClearColor(*self.bg_color, 1.0)
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
        glLoadIdentity()


        # 相机位置
        gluLookAt(0, 0, self.zoom,
                  0, 0, 0,
                  0, 1, 0)

        # 应用旋转
        glRotatef(self.rotation_x, 1, 0, 0)
        glRotatef(self.rotation_y, 0, 1, 0)

        # 显示坐标轴
        if self.show_axis:
            self.draw_coordinate_system()


        # 绘制点云（修复透明度）
        if self.point_cloud is not None:
            glDisable(GL_LIGHTING)  # 禁用光照，确保透明度生效
            glEnable(GL_BLEND)  # 启用混合模式（关键）
            glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA)  # 设置混合因子
            glPointSize(self.point_size)
            glBegin(GL_POINTS)
            for i, point in enumerate(self.point_cloud):
                if self.point_colors is not None and i < len(self.point_colors):
                    color = self.point_colors[i]
                    glColor4f(color[0], color[1], color[2], self.opacity)  # 透明度生效
                else:
                    glColor4f(0.4, 0.7, 1.0, self.opacity)  # 默认蓝+透明度
                glVertex3f(point[0], point[1], point[2])
            glEnd()
            glDisable(GL_BLEND)  # 禁用混合
            glEnable(GL_LIGHTING)  # 恢复光照

        # 绘制网格
        if self.mesh is not None:
            vertices = np.asarray(self.mesh.vertices)
            triangles = np.asarray(self.mesh.triangles)
            normals = np.asarray(self.mesh.vertex_normals)

            # 中心化和缩放
            centroid = np.mean(vertices, axis=0)
            vertices = vertices - centroid
            max_range = np.max(np.linalg.norm(vertices, axis=1))
            if max_range > 0:
                vertices = vertices / max_range * 2.0

            # 网格材质（适配黑色背景）
            glMaterialfv(GL_FRONT_AND_BACK, GL_AMBIENT, [0.3, 0.7, 1.0, 1.0])
            glMaterialfv(GL_FRONT_AND_BACK, GL_DIFFUSE, [0.6, 1.0, 1.0, 1.0])
            glMaterialfv(GL_FRONT_AND_BACK, GL_SPECULAR, [1.0, 1.0, 1.0, 1.0])
            glMaterialf(GL_FRONT_AND_BACK, GL_SHININESS, 30.0)

            # 绘制三角形网格
            glBegin(GL_TRIANGLES)
            for triangle in triangles:
                for idx in triangle:
                    vertex = vertices[idx]
                    normal = normals[idx] if len(normals) > idx else [0, 0, 1]
                    normal = normal / np.linalg.norm(normal) if np.linalg.norm(normal) > 0 else [0, 0, 1]
                    glNormal3f(normal[0], normal[1], normal[2])
                    glVertex3f(vertex[0], vertex[1], vertex[2])
            glEnd()

            # 绘制线框
            glColor3f(0.2, 0.4, 0.6)
            glLineWidth(0.8)
            glDisable(GL_LIGHTING)
            glBegin(GL_LINES)
            for triangle in triangles:
                for i in range(3):
                    v1 = vertices[triangle[i]]
                    v2 = vertices[triangle[(i+1)%3]]
                    glVertex3f(v1[0], v1[1], v1[2])
                    glVertex3f(v2[0], v2[1], v2[2])
            glEnd()
            glEnable(GL_LIGHTING)

    def draw_coordinate_system(self):
        """绘制坐标轴（X红、Y绿、Z蓝）"""
        glDisable(GL_LIGHTING)
        glLineWidth(2.0)
        # X轴（红色）
        glColor3f(1.0, 0.0, 0.0)
        glBegin(GL_LINES)
        glVertex3f(0, 0, 0)
        glVertex3f(1.0, 0, 0)
        glEnd()
        # Y轴（绿色）
        glColor3f(0.0, 1.0, 0.0)
        glBegin(GL_LINES)
        glVertex3f(0, 0, 0)
        glVertex3f(0, 1.0, 0)
        glEnd()
        # Z轴（蓝色）
        glColor3f(0.0, 0.0, 1.0)
        glBegin(GL_LINES)
        glVertex3f(0, 0, 0)
        glVertex3f(0, 0, 1.0)
        glEnd()
        glEnable(GL_LIGHTING)

    # ---------------------- 新增：属性更新接口 ----------------------
    def set_point_cloud(self, point_cloud_data, point_colors=None, point_size=None, opacity=None, show_axis=None):
        """设置点云并可选指定属性"""
        self.clear()
        if point_cloud_data is not None:
            # 中心化和缩放
            centroid = np.mean(point_cloud_data, axis=0)
            point_cloud_data = point_cloud_data - centroid
            max_range = np.max(np.linalg.norm(point_cloud_data, axis=1))
            if max_range > 0:
                point_cloud_data = point_cloud_data / max_range * 2.0

        self.point_cloud = point_cloud_data
        self.point_colors = point_colors
        # 可选更新属性
        if point_size is not None:
            self.point_size = point_size
        if opacity is not None:
            self.opacity = opacity
        if show_axis is not None:
            self.show_axis = show_axis
        self.update()

    def update_point_size(self, size):
        """更新点大小"""
        self.point_size = size
        self.update()

    def update_opacity(self, opacity):
        """更新透明度"""
        self.opacity = opacity
        self.update()

    def update_bg_color(self, r, g, b):
        """更新背景色"""
        self.bg_color = (r, g, b)
        glClearColor(r, g, b, 1.0)
        self.update()

    def update_show_axis(self, show):
        """更新坐标轴显示状态"""
        self.show_axis = show
        self.update()

    def set_mesh(self, mesh):
        """设置网格（保留原有）"""
        self.clear()
        self.mesh = mesh
        self.rotation_x = 0
        self.rotation_y = 0
        self.zoom = 5.0
        self.update()

    def clear(self):
        """清空（保留原有）"""
        self.point_cloud = None
        self.point_colors = None
        self.mesh = None
        self.update()

    # ---------------------- 鼠标交互（保留原有）----------------------
    def mousePressEvent(self, event):
        self.last_mouse_pos = event.pos()

    def mouseMoveEvent(self, event):
        if self.last_mouse_pos is not None:
            dx = event.x() - self.last_mouse_pos.x()
            dy = event.y() - self.last_mouse_pos.y()
            self.rotation_y += dx * 0.5
            self.rotation_x += dy * 0.5
            self.last_mouse_pos = event.pos()
            self.update()

    def wheelEvent(self, event):
        delta = event.angleDelta().y()
        self.zoom -= delta * 0.001
        self.zoom = max(1.0, min(self.zoom, 20.0))
        self.update()